'''
数据集预处理脚本 messidor
分类数据集
'''
import sys, os
import csv
import json
from data_preprocess.utils.crop import crop_resize_save
from data_preprocess.utils.stat import save_classification_stats
from tqdm import tqdm
def parse_csv(csv_path):
    """
    读取 csv 文件，并返回一个字典，格式如下：
    
    {
       image_id: "<diagnosis text>"
    }
    
    其中，文本描述的生成规则为：
      - 对于 adjudicated_dr_grade (ICDR 5 级):
           0 -> normal
           1 -> Mild Diabetic Retinopathy
           2 -> Moderate Diabetic Retinopathy
           3 -> Severe Diabetic Retinopathy
           4 -> Proliferative Diabetic Retinopathy
      - 对于 adjudicated_dme:
           0 -> normal
           1 -> Referable Diabetic Macular Edema
      - 如果两者均为 normal，则文本为 "normal"
      - 如果只有一项异常，则只写异常项的描述
      - 如果两项都异常，则返回 "【DR condition】, 【DME condition】"（先 DR 后 DME）
      
      对于 adjudicated_gradable 为 0 的记录，不生成文本（或可生成空字符串）
    """
    # 定义映射表
    dr_mapping = {
        "0": "normal",
        "1": "Mild Diabetic Retinopathy",
        "2": "Moderate Diabetic Retinopathy",
        "3": "Severe Diabetic Retinopathy",
        "4": "Proliferative Diabetic Retinopathy"
    }
    
    dme_mapping = {
        "0": "normal",
        "1": "Referable Diabetic Macular Edema"
    }
    
    quality_mapping={
        "0":"Ungradable",
        "1":"Gradable"
    }
    result = {}
    with open(csv_path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            image_id = row.get("image_id", "").strip()
            # 若图像不可评分，则跳过该记录
            gradable = row.get("adjudicated_gradable", "").strip()
            quality = quality_mapping[gradable]
        
            dr_val = row.get("adjudicated_dr_grade", "").strip()
            dme_val = row.get("adjudicated_dme", "").strip()
            
            # 若任一值为空，跳过
            if dr_val == "" or dme_val == "":
                if quality != "Ungradable":
                    raise ValueError(f"缺少必要字段，行数据：{row}")
                else:
                    result[image_id] = {
                        "quality":quality,
                    }
                continue
            # 获取映射结果
            dr_text = dr_mapping.get(dr_val, dr_val)
            dme_text = dme_mapping.get(dme_val, dme_val)
            
            # 根据规则生成描述文本
            if dr_text == "normal" and dme_text == "normal":
                full_text = "normal"
            elif dr_text == "normal" and dme_text != "normal":
                full_text = dme_text
            elif dr_text != "normal" and dme_text == "normal":
                full_text = dr_text
            else:
                full_text = f"{dr_text}, {dme_text}"
            
            result[image_id] = {
                "text": full_text,
                "quality":quality
            }
    return result

def gather_data(data_path, tar_path, prefix='messidor_', resize=(512,512)):
    """
    对新数据集进行预处理：
      - 数据集根目录下要求存在 IMAGES 文件夹和 messidor_label.csv 文件；
      - 将所有处理后的图像统一保存到 tar_path/images 下，
        新图像名称格式为 {prefix}_{原始文件名}.png（保留原始文件名，仅修改后缀并加上前缀）。
      - 从 CSV 中解析标注信息，并将每个图像对应的标注写入 data_dict，
        格式为：
            {
                new_image_name: {
                    "image_name": new_image_name,
                    "image_path": "images/{new_image_name}",
                    "diagnosis": {
                        "classification": {
                        "text": <diagnosis text>
                        }
                    }
                },
                ...
            }
      - 最后将 data_dict 保存到 tar_path/annotations.json 中。
    
    参数：
        data_path (str): 新数据集根目录，要求其下包含 IMAGES 文件夹和 messidor_label.csv 文件；
        tar_path (str): 预处理后数据存放目录；
        prefix (str): 处理后图片名称的前缀，默认为 "messidor_"
        resize (tuple): 目标尺寸，默认为 (512,512)
    
    返回：
        dict: 标注信息字典
    """
    # 创建目标目录
    os.makedirs(tar_path, exist_ok=True)
    images_dir = os.path.join(tar_path, "images")
    os.makedirs(images_dir, exist_ok=True)
    
    # 解析 CSV 获取标注信息
    label_dict = parse_csv(os.path.join(data_path, "messidor_label.csv"))
    
    data_dict = {}
    image_src_path = os.path.join(data_path, "IMAGES")
    image_name_list = sorted(os.listdir(image_src_path))
    
    for image_file in tqdm(image_name_list,desc='process image',unit='images'):
        if not image_file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.ppm')):
            continue

        base_name = os.path.splitext(image_file)[0]
        # 新的图像名称：保留原始文件名，仅加前缀并统一保存为 PNG 格式
        new_image_name = f"{prefix}_{base_name}.png"
        dest_image_path = os.path.join(images_dir, new_image_name)
        src_image_path = os.path.join(image_src_path, image_file)
        # 调用 crop_resize_save 进行裁剪与 resize
        crop_info=crop_resize_save(
            image_path=src_image_path,
            save_path=dest_image_path,
            
            resize=resize,
            crop_threshold=15
        )
        data_dict[new_image_name] = {
            "image_name": new_image_name,
            "image_path": dest_image_path,
            'original_path': src_image_path,
            'crop_info':crop_info,
            
        }
        # 获取标注文本，若未找到则默认为空字符串,就没有diagnosis字段
        if image_file[-3:]=="JPG":
            image_file=image_file[:-3]+'jpg'# 原数据集存储jpg的时候用的是大写，和标签csv的小写冲突，理想实现应该label_dict在构造的时候全用小写，这里也全处理成小写，但简明处理jpg依旧是一种方法

        diag_text = label_dict.get(image_file, "")
        
        if diag_text:
            data_dict[new_image_name]["diagnosis"]={
                "classification": diag_text
            }
        else:
            print(f"{image_file} not in label_dict")
            print(list(label_dict.keys())[:10])
            raise

    # 保存统计
    save_classification_stats(data_dict, './experiments/stat/messidor.json')
    # 保存标注信息到 tar_path/annotations.json
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(data_dict, f, indent=4)
    
    return data_dict

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Messidor 数据集预处理")
    parser.add_argument("--data_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_dataset/Messidor",
                        help="原始数据集根目录，要求其下包含 IMAGES 文件夹和 messidor_label.csv")
    parser.add_argument("--tar_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_processed/Messidor",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="messidor",
                        help="处理后图片名称的前缀，默认为 'messidor'")
    parser.add_argument("--resize", type=int, nargs=2, default=[512,512],
                        help="目标尺寸，默认为 512 512")
    args = parser.parse_args()
    
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix, resize=tuple(args.resize))
    print("Preprocessing completed. Annotations saved.")
